from django.db import migrations
import random

def insert_word_trends(apps, schema_editor):
    WordTrend = apps.get_model('trend', 'WordTrend')
    
    # 定义所有核心词
    keywords = [
        '人工智能', '机器学习', '深度学习', '神经网络', '计算机视觉', 
        '自然语言处理', '大数据', '云计算', '区块链', '物联网',
        '5G通信', '边缘计算', '量子计算', '增强现实', '虚拟现实',
        '自动驾驶', '机器人', '数字孪生', '网络安全', '数据挖掘',
        '图像识别', '语音识别', '推荐系统', '知识图谱', '强化学习',
        '迁移学习', '联邦学习', '生成对抗网络', '智能制造', '智慧城市'
    ]
    
    # 为每个关键词生成基础数量(确保某些词总是较多)
    base_counts = {
        keyword: random.randint(100, 2000) 
        for keyword in keywords
    }
    
    # 特别设置一些热门词的基础量更大
    hot_keywords = ['人工智能', '机器学习', '深度学习', '大数据', '云计算']
    for keyword in hot_keywords:
        base_counts[keyword] = random.randint(1800, 3200)
    
    trends_data = []
    
    # 生成2019-2023年的数据
    for year in range(2019, 2024):
        for keyword in keywords:
            base = base_counts[keyword]
            # 添加随机波动，可能增长也可能下降
            variation = random.uniform(-0.3, 0.5)  # -30% 到 +50% 的波动
            patent_count = int(base * (1 + variation))
            # 确保专利数量不会太小
            patent_count = max(50, patent_count)
            # 词出现次数略多于专利数量
            count = int(patent_count * random.uniform(1.2, 1.8))
            
            trends_data.append({
                'year': year,
                'word': keyword,
                'patent_count': patent_count,
                'count': count
            })
    
    # 批量创建数据
    WordTrend.objects.bulk_create([
        WordTrend(**data) for data in trends_data
    ])

def reverse_word_trends(apps, schema_editor):
    WordTrend = apps.get_model('trend', 'WordTrend')
    WordTrend.objects.all().delete()

class Migration(migrations.Migration):
    dependencies = [
        ('trend', '05_ins_top_orgs'),
    ]

    operations = [
        migrations.RunPython(insert_word_trends, reverse_word_trends),
    ] 